from pathlib import Path
from utils.tools import *
import logging
from Matchers.superglue.superglue import SuperGlue

from DataLoader.SequenceImageLoader import SequenceImageLoader
from DataLoader.KITTILoader import KITTILoader
from Detectors.SuperPointDetector import SuperPointDetector

from models.utils import AverageTimer


import cv2


class SuperGlueMatcher(object):
    default_config = {
        "descriptor_dim": 256,
        "weights": "outdoor",
        "keypoint_encoder": [32, 64, 128, 256],
        "GNN_layers": ["self", "cross"] * 9,
        "sinkhorn_iterations": 100,
        "match_threshold": 0.2,
        "cuda": True
    }

    def __init__(self, config={}):
        self.config = self.default_config
        self.config = {**self.config, **config}
        logging.info("SuperGlue matcher config: ")
        logging.info(self.config)

        # self.device = 'cuda' if torch.cuda.is_available() and isinstance(self.config["cuda"], int)else 'cpu'
        # print(f"device: {self.device}")
        if("device" in config.keys()):
            self.device = self.config["device"]
        else:
            self.device = 'cuda' if torch.cuda.is_available() and isinstance(self.config["cuda"], int)else 'cpu'


        assert self.config['weights'] in ['indoor', 'outdoor']
        path = Path(__file__).parent
        path = path / 'superglue/superglue_{}.pth'.format(self.config['weights'])
        self.config["path"] = path

        logging.info("creating SuperGlue matcher...")
        self.superglue = SuperGlue(self.config).to(self.device)


    def __call__(self, pred):
        # setup data for superglue
        logging.debug("prepare input data for superglue...")
        # data={}
        # data['image_size0'] = torch.from_numpy(kptdescs["ref"]["image_size"]).float().to(self.device)
        # data['image_size1'] = torch.from_numpy(kptdescs["cur"]["image_size"]).float().to(self.device)

        # data['scores0'] = kptdescs["ref"]["scores"][0].unsqueeze(0)
        # data['keypoints0'] = kptdescs["ref"]["keypoints"][0].unsqueeze(0)
        # data['descriptors0'] = kptdescs["ref"]["descriptors"][0].unsqueeze(0)
        # data['scores1'] = kptdescs["cur"]["scores"][0].unsqueeze(0)
        # data['keypoints1'] = kptdescs["cur"]["keypoints"][0].unsqueeze(0)
        # data['descriptors1'] = kptdescs["cur"]["descriptors"][0].unsqueeze(0)
        # data['ref']['image'] = imgs["ref"]
        # data['cur']['image'] = imgs['cur']

        # Forward !!
        logging.debug("matching keypoints with superglue...")
        # data = self.superglue(data)
        pred = {**pred, **self.superglue(pred)}

        # get matching keypoints
        kpts0 = pred["ref"]["keypoints"]
        kpts1 = pred["cur"]["keypoints"]

        # matches = pred['matches0'][0].cpu().numpy()
        # confidence = pred['matching_scores0'][0].cpu().detach().numpy()

        # # Sort them in the order of their confidence.
        # match_conf = []
        # for i, (m, c) in enumerate(zip(matches, confidence)):
        #     match_conf.append([i, m, c])
        # match_conf = sorted(match_conf, key=lambda x: x[2], reverse=True)
        #
        # valid = [[l[0], l[1]] for l in match_conf if l[1] > -1]
        # v0 = [l[0] for l in valid]
        # v1 = [l[1] for l in valid]
        # mkpts0 = kpts0[v0]
        # mkpts1 = kpts1[v1]

        # Keep the matching keypoints.
        valid = pred['matches0'][0] > -1
        mkpts0 = kpts0[valid]
        mkpts1 = kpts1[pred['matches0'][0][valid]]
        mconf = pred['matching_scores0'][0].cpu().detach()[valid]

        ret_dict = {
            "ref_keypoints": mkpts0,
            "cur_keypoints": mkpts1,
            # "match_score": confidence[v0]
            "match_score": mconf
        }

        return ret_dict


if __name__ == "__main__":
    torch.set_grad_enabled(False)

    loader = SequenceImageLoader()
    detector = SuperPointDetector({"cuda": 0})
    matcher = SuperGlueMatcher({"cuda": 0, "weights": "outdoor"})

    kptdescs = {'cur': None, 'ref': None}
    imgs = {}
    timer = AverageTimer(newline=True)

    # for i, img in enumerate(loader):
    #     # imgs["cur"] = torch.from_numpy(img/255.).float()[None, None].to(matcher.device)
    #     imgs["cur"] = img
    #     kptdescs["cur"] = detector(img)
    #     if i >= 1:
    #         matches = matcher(kptdescs)
    #         timer.update('matcher')
    #
    #         img = plot_matches(imgs['ref'], imgs['cur'],
    #                            matches['ref_keypoints'].cpu().detach().numpy()[0::5], matches['cur_keypoints'].cpu().detach().numpy()[0::5],
    #                            matches['match_score'].cpu().detach().numpy()[0::5], layout='lr')
    #         cv2.imshow("track", img)
    #         # if cv2.waitKey(1) == 27:     # 按 ‘ESC'结束运行
    #         #     break
    #
    #         cv2.waitKey(1)
    #         # matches.clear()
    #         # torch.cuda.empty_cache()
    #
    #     kptdescs["ref"], imgs["ref"] = kptdescs["cur"], imgs["cur"]
    #     # kptdescs["ref"] = kptdescs["cur"]


    imgs["cur"] = cv2.imread("/home/daybeha/Documents/Dataset/Kitti/sequences/05/image_0/000014.png")
    imgs["ref"] = cv2.imread("/home/daybeha/Documents/Dataset/Kitti/sequences/05/image_0/002430.png")
    kptdescs["cur"] = detector(imgs["cur"])
    kptdescs["ref"] = detector(imgs["ref"])
    matches = matcher(kptdescs)
    score = np.mean(matches["match_score"].cpu().detach().numpy())
    num = matches["match_score"].shape[0]
    print(f"score: {score}   num: {num}")


    timer.update('matcher')
    img = plot_all(imgs['ref'], imgs['cur'],
                       matches['ref_keypoints'].cpu().detach().numpy()[0::5], matches['cur_keypoints'].cpu().detach().numpy()[0::5],
                       matches['match_score'].cpu().detach().numpy()[0::5])
    cv2.imshow("track", img)
    # if cv2.waitKey(1) == 27:     # 按 ‘ESC'结束运行
    #     break

    cv2.waitKey()

